import torch
import numpy as np
from PIL import Image

def simple_call():
    from diffusers import DDPMPipeline
    ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")
    image = ddpm(num_inference_steps=25).images[0]
    return image

def detail_call():
    from diffusers import DDPMScheduler, UNet2DModel
    #定义调度器
    scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
    #设置调度器的时间步长
    scheduler.set_timesteps(50)
    #定义模型
    model = UNet2DModel.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")
    #随机生成一张噪声图像
    sample_size = model.config.sample_size
    noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")
    #循环降噪过程
    input = noise
    for t in scheduler.timesteps:
        with torch.no_grad():
            #生一个噪声残差作为下一步的输入
            noisy_residual = model(input, t).sample
        previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
        input = previous_noisy_sample
    #将最终的噪声残差转换为图像
    image = (input / 2 + 0.5).clamp(0, 1).squeeze()
    image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy()
    image = Image.fromarray(image)
    return image

#image = simple_call()
image = detail_call()
image.save("tmp/c01.png")